package com.simple.rpc.common.network;

import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 项目: simple-rpc
 *
 * 功能描述: 线程池包装
 *
 * @author: WuChengXing
 *
 * @create: 2023-08-18 17:17
 **/
public class SimpleExecutorWrapper extends ThreadPoolExecutor {

    private final ExecutorService executor;

    private SimpleExecutorWrapper(ThreadPoolExecutor executor) {
        super(executor.getCorePoolSize(), executor.getMaximumPoolSize(), executor.getKeepAliveTime(TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS, executor.getQueue(), executor.getThreadFactory(), executor.getRejectedExecutionHandler());
        this.executor = executor;
    }

    /**
     * 创建一个新的executor
     *
     * @param executor
     *
     * @return
     */
    public static SimpleExecutorWrapper createByThreadPoolExecutor(ThreadPoolExecutor executor) {
        return new SimpleExecutorWrapper(executor);
    }

    @Override
    public void execute(Runnable command) {
        this.executor.execute(wrapperRunnable(command));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return executor.submit(task);
    }

    /**
     * 包装
     *
     * @param command
     *
     * @return
     */
    private Runnable wrapperRunnable(Runnable command) {
        return new SimpleRunnableWrapper(command);
    }

    /**
     * runnable包装类
     */
    class SimpleRunnableWrapper implements Runnable {

        private final AtomicReference<Map<SimpleThreadLocal<?>, Object>> tempSaveThreadLocal = new AtomicReference<>(SimpleThreadLocal.copy());

        private final Runnable runnable;

        private SimpleRunnableWrapper(Runnable runnable) {
            this.runnable = runnable;
        }


        @Override
        public void run() {
            SimpleThreadLocal.initChildThreadLocal(tempSaveThreadLocal.get());
            try {
                runnable.run();
            } finally {
                SimpleThreadLocal.removeChildThreadLocal();
                tempSaveThreadLocal.set(null);
            }
        }
    }

}
